from python_sklearn.dataset import mydataset
import numpy as np
import pandas as pd
import csv
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    data = mydataset().getData()
    x = pd.DataFrame()
    for col in ['user_id', 'cat_id', 'cat1', 'gender', 'p']:
        t = pd.get_dummies(data[col], prefix=col)
        x = pd.concat((x, t), axis=1)
    print(x)
    x=x.groupby(by=x.index,level=None).sum()
    # 将求和的转为1化为one-hot
    x[x!=0]=1


    x1=x
    print(x)
    # 如果预测是非数字需要编码化
    # y = np.array(pd.Categorical(data['live']).codes)
    y=data['live']
    y1=y
    print(y)
    y=y.groupby(by=y.index,level=None).mean()
    print(y)


    x, x_test, y, y_test = train_test_split(x, y, test_size=0.3)
    clf = RandomForestRegressor(n_estimators=60, max_depth=7)
    clf.fit(x, y)

    y_test_hat = clf.predict(x_test)
    y_hat = clf.predict(x)
    # print(y_test_hat)
    # print(clf.score(x_test, y_test_hat))
    x1.to_csv("out.csv",index=True)


    a=0
    for i in range(len(y_test_hat)):
        if -100<y_test.tolist()[i]-y_test_hat.tolist()[i]<100:
            a=a+1
    print(a/len(y_test_hat))
    with open('oben.csv', 'w') as f:
        writer = csv.writer(f)
        writer.writerow(y)
        writer.writerow(y1)
        # writer.writerow(y)
        # writer.writerow(y)
        # writer.writerow(y)
        # writer.writerow(y_test)
        # writer.writerow(y_test_hat)
        # writer.writerow(y_hat)
